import torch
from torch import nn

class TestNet(nn.Module):
    def __init__(self):
        super(TestNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, dilation=2)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1)
    def forward(self, x):
        x = self.conv1(x)
        x = self.pool1(x)
        return x


model = TestNet()
print(model)
input_tensor = torch.randn(1, 3, 1080, 1920)  # 假设输入图像为1080P
output_tensor = model(input_tensor)
print(output_tensor.shape)

# conv_layer = LayerInfo(
#     layer_id='conv1',
#     layer_type='conv2d',
#     kernel_size=(3,3),
#     stride=(2,2),
#     padding=(1,1),
#     dilation=(2,2),
#     input_shape=(1,1080,1920,3),
#     output_shape=(1,540,960,64)
# )
# rd.register_layer(conv_layer)

